#!/usr/bin/env python3
"""
STM32 Firmware Side-Channel AttackTrace采集器（简化版）
SupportAES、SM4、SM2、RSA、ECC、Dilithiumetc.Algorithm
每种Algorithm固定一种操作
"""

import subprocess
import os
import time
import argparse
import numpy as np
import psutil
import threading
from typing import List, Dict, Optional
from dataclasses import dataclass
import re
from pathlib import Path

@dataclass
class Trace:
    """单条Trace窗口（按TRACE_START/TRACE_END）"""
    index: int
    start_time: float
    end_time: float
    duration: float
    cpu_usage: List[float]  # CPUUse率Sampling

class CPUMonitor:
    """QEMUProcessCPUMonitor器（超高频Sampling）"""
    def __init__(self, sample_interval: float = 0.0001):  # 0.1ms
        self.sample_interval = sample_interval
        self.monitoring = False
        self.cpu_samples = []
        self.monitor_thread = None
        self._ready = False
        self.qemu_process = None
    
    def _monitor_loop(self):
        """MonitorLoop（超高频率）"""
        import time as time_module
        self._ready = True
        
        while True:
            if self.monitoring and self.qemu_process:
                try:
                    cpu_percent = self.qemu_process.cpu_percent(interval=None)
                    self.cpu_samples.append(cpu_percent)
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    break
            else:
                time_module.sleep(0.001)
    
    def set_qemu_process(self, qemu_process):
        """Settings要Monitor的QEMUProcess"""
        self.qemu_process = psutil.Process(qemu_process.pid)
        self.qemu_process.cpu_percent(interval=None)
    
    def start_monitoring(self):
        """StartMonitor"""
        if self.monitor_thread is None:
            self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
            self.monitor_thread.start()
            while not self._ready:
                time.sleep(0.0001)
        
        self.cpu_samples = []
        self.monitoring = True
    
    def stop_monitoring(self) -> List[float]:
        """停止Monitor并ReturnSamplingData"""
        self.monitoring = False
        time.sleep(0.001)
        return self.cpu_samples.copy()

class FirmwareTraceCollector:
    """通用固件Trace采集器"""
    
    # AlgorithmConfiguration（简化版：每种Algorithm固定一种操作）
    ALGORITHM_CONFIGS = {
        'aes': {
            'name': 'AES',
            'key_size': 16,
            'block_size': 16,
            'operation': 'encrypt',  # 固定For加密
            'description': 'AES-128加密'
        },
        'sm4': {
            'name': 'SM4',
            'key_size': 16,
            'block_size': 16,
            'operation': 'encrypt',  # 固定For加密
            'description': 'SM4加密'
        },
        'sm2': {
            'name': 'SM2',
            'key_size': 32,
            'block_size': 32,
            'operation': 'sign',  # 固定For签名
            'description': 'SM2Digital Signature'
        },
        'rsa': {
            'name': 'RSA',
            'key_size': 256,
            'block_size': 256,
            'operation': 'decrypt',  # 固定For解密
            'description': 'RSA私钥解密'
        },
        'ecc': {
            'name': 'ECC',
            'key_size': 32,
            'block_size': 32,
            'operation': 'sign',  # 固定For签名
            'description': 'ECCDigital Signature(P-256)'
        },
        'dilithium': {
            'name': 'Dilithium',
            'key_size': 32,
            'block_size': 32,
            'operation': 'sign',  # 固定For签名
            'description': 'DilithiumDigital Signature'
        }
    }
    
    def __init__(self, algorithm: str = 'aes', work_dir: Optional[str] = None):
        """
        Initialize采集器
        
        Args:
            algorithm: Algorithm名称 (aes, sm4, sm2, rsa, ecc, dilithium)
            work_dir: Working directory，DefaultFor stm32_{algorithm}/
        """
        self.algorithm = algorithm.lower()
        if self.algorithm not in self.ALGORITHM_CONFIGS:
            raise ValueError(f"NotSupport的Algorithm: {algorithm}")
        
        self.config = self.ALGORITHM_CONFIGS[self.algorithm]
        self.work_dir = work_dir or f"stm32_{self.algorithm}"
        
        # CheckWorking directory
        if not os.path.exists(self.work_dir):
            raise ValueError(f"Working directoryNot存In: {self.work_dir}")
        
        self.traces = []
        self.cpu_monitor = CPUMonitor()
        
        print(f"Initialize {self.config['name']} Trace采集器")
        print(f"Working directory: {self.work_dir}")
    
    def generate_test_data(self, num_traces: int, key: Optional[str] = None, 
                          delay_us: int = 1000):
        """GenerateTestDataHeader file（UseAlgorithm固定的操作）"""
        # Use父Directory的generate_test_data.py
        script = os.path.join(os.path.dirname(self.work_dir), 'generate_test_data.py')
        if not os.path.exists(script):
            print(f"Warning: {script} Not存In，SkipTestDataGenerate")
            return True
        
        operation = self.config['operation']
        description = self.config['description']
        
        print(f"\n[1/4] GenerateTestData...")
        print(f"  Algorithm: {description}")
        print(f"  操作Class型: {operation} (固定)")
        print(f"  Trace数量: {num_traces}")
        print(f"  延时: {delay_us}微秒")
        
        cmd = [
            'python3', script,
            '--algorithm', self.algorithm,
            '--traces', str(num_traces),
            '--delay', str(delay_us),
            '--output', os.path.join(self.work_dir, f'{self.algorithm}_test_data.h')
        ]
        
        if key:
            cmd.extend(['--key', key])
        
        try:
            result = subprocess.run(cmd, check=True, 
                                  capture_output=True, text=True)
            print(result.stdout)
            return True
        except subprocess.CalledProcessError as e:
            print(f"Error: GenerateTestDataFailed")
            print(e.stderr)
            return False
    
    def _compile_firmware(self) -> bool:
        """编译固件"""
        print(f"\n[2/4] 编译固件...")
        
        makefile = os.path.join(self.work_dir, 'Makefile')
        if not os.path.exists(makefile):
            print(f"Error: Makefile Not存In: {makefile}")
            return False
        
        try:
            result = subprocess.run(
                ['make', 'clean'],
                cwd=self.work_dir,
                check=False,
                capture_output=True
            )
            
            result = subprocess.run(
                ['make', '-j4'],
                cwd=self.work_dir,
                check=True,
                capture_output=True,
                text=True
            )
            
            print("[OK] 固件编译Success")
            
            # FindGenerate的.elfFile
            build_dir = os.path.join(self.work_dir, 'build')
            if os.path.exists(build_dir):
                elf_files = list(Path(build_dir).glob('*.elf'))
                if elf_files:
                    print(f"  固件: {elf_files[0].name}")
            
            return True
            
        except subprocess.CalledProcessError as e:
            print(f"Error: 编译Failed")
            print(e.stderr)
            return False
    
    def run_qemu_and_collect(self):
        """RunQEMU并采集Trace"""
        print(f"\n[3/4] 启动QEMU并采集Trace...")
        
        # Find.elfFile
        build_dir = os.path.join(self.work_dir, 'build')
        elf_files = list(Path(build_dir).glob('*.elf'))
        if not elf_files:
            print(f"Error: 未找To.elfFileIn {build_dir}")
            return False
        
        elf_file = str(elf_files[0])
        print(f"  固件: {os.path.basename(elf_file)}")
        
        # 启动QEMU
        qemu_cmd = [
            'qemu-system-arm',
            '-M', 'netduinoplus2',
            '-nographic',
            '-semihosting-config', 'enable=on,target=native',
            '-kernel', elf_file
        ]
        
        try:
            qemu = subprocess.Popen(
                qemu_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1
            )
            
            # SettingsQEMUProcess
            self.cpu_monitor.set_qemu_process(qemu)
            
            # ReadOutput并采集Trace
            current_trace = None
            
            for line in qemu.stdout:
                line = line.strip()
                
                if line:
                    print(line)
                
                # ParseTRACE_START
                if line.startswith('TRACE_START:'):
                    index = int(line.split(':')[1])
                    self.cpu_monitor.start_monitoring()
                    start_time = time.time()
                    current_trace = {
                        'index': index,
                        'start_time': start_time
                    }
                
                # ParseTRACE_END
                elif line.startswith('TRACE_END:') and current_trace:
                    cpu_usage = self.cpu_monitor.stop_monitoring()
                    end_time = time.time()
                    
                    trace = Trace(
                        index=current_trace['index'],
                        start_time=current_trace['start_time'],
                        end_time=end_time,
                        duration=end_time - current_trace['start_time'],
                        cpu_usage=cpu_usage
                    )
                    
                    self.traces.append(trace)
                    current_trace = None
                
                # 批ProcessCompleted
                elif line == 'BATCH_COMPLETE':
                    print(f"\n[OK] 批ProcessCompleted")
                    break
        
        except KeyboardInterrupt:
            print("\nInterrupt采集")
        finally:
            qemu.terminate()
            qemu.wait()
        
        print(f"\n采集Completed: {len(self.traces)} 条Trace")
        return len(self.traces) > 0
    
    def save_traces(self, output_file: str = "traces.npz", 
                   include_key: bool = True,
                   include_plaintexts: bool = False):
        """SaveTraceData"""
        if not self.traces:
            print("没HasTraceData可Save")
            return False
        
        print(f"\n[4/4] SaveTraceData...")
        
        # 准备基础Data
        indices = np.array([t.index for t in self.traces], dtype=np.int32)
        start_times = np.array([t.start_time for t in self.traces], dtype=np.float64)
        end_times = np.array([t.end_time for t in self.traces], dtype=np.float64)
        durations = np.array([t.duration for t in self.traces], dtype=np.float32)
        
        # CPUTraceData
        max_len = max(len(t.cpu_usage) for t in self.traces) if self.traces else 0
        cpu_traces = np.zeros((len(self.traces), max_len), dtype=np.float32)
        for i, t in enumerate(self.traces):
            cpu_len = len(t.cpu_usage)
            if cpu_len > 0:
                cpu_traces[i, :cpu_len] = t.cpu_usage
        
        # 准备SaveData
        save_data = {
            'indices': indices,
            'start_times': start_times,
            'end_times': end_times,
            'durations': durations,
            'cpu_traces': cpu_traces,
            'traces': cpu_traces,  # 兼容性
            'algorithm': self.algorithm,
            'operation': self.config['operation'],  # Use固定操作
            'num_traces': len(self.traces)
        }
        
        # FromTestDataHeader fileReadKey（IfRequire）
        if include_key:
            key = self._load_key_from_header()
            if key is not None:
                save_data['keys'] = key  # Use'keys'以MatchSide-Channel AttackModule的期望
        
        # FromTestDataHeader fileReadPlaintext（IfRequire）
        if include_plaintexts:
            plaintexts = self._load_plaintexts_from_header(len(self.traces))
            if plaintexts is not None:
                save_data['plaintexts'] = plaintexts
        
        # Save
        np.savez_compressed(output_file, **save_data)
        
        print(f"[OK] TraceData已Save: {output_file}")
        print(f"  DataInclude:")
        for key, value in save_data.items():
            if isinstance(value, np.ndarray):
                print(f"    - {key}: {value.shape} {value.dtype}")
            else:
                print(f"    - {key}: {value}")
        
        return True
    
    def _load_key_from_header(self) -> Optional[np.ndarray]:
        """FromTestDataHeader fileLoadKey"""
        header_file = os.path.join(self.work_dir, f"{self.algorithm}_test_data.h")
        if not os.path.exists(header_file):
            return None
        
        try:
            with open(header_file, 'r') as f:
                content = f.read()
            
            # FindKey定义（SupportArray大小For宏定义）
            pattern = rf'{self.algorithm}_key\[[^\]]+\]\s*=\s*\{{([^}}]+)\}}'
            match = re.search(pattern, content, re.IGNORECASE)
            
            if match:
                hex_values = re.findall(r'0x([0-9a-fA-F]{2})', match.group(1))
                key = np.array([int(h, 16) for h in hex_values], dtype=np.uint8)
                print(f"  FromHeader fileLoadKey: {len(key)} 字节")
                return key
        except Exception as e:
            print(f"  Warning: LoadKeyFailed: {e}")
        
        return None
    
    def _load_plaintexts_from_header(self, num_traces: int) -> Optional[np.ndarray]:
        """FromTestDataHeader fileLoadPlaintext"""
        header_file = os.path.join(self.work_dir, f"{self.algorithm}_test_data.h")
        if not os.path.exists(header_file):
            return None
        
        try:
            with open(header_file, 'r') as f:
                content = f.read()
            
            # FindTestDataArray（Support宏定义的Array大小）
            pattern = r'test_data\[[^\]]+\]\[[^\]]+\]\s*=\s*\{(.+)\}'
            match = re.search(pattern, content, re.DOTALL)
            
            if match:
                hex_values = re.findall(r'0x([0-9a-fA-F]{2})', match.group(1))
                
                # TryFromConfiguration中获取块大小
                block_size = self.config.get('block_size', 16)
                
                if len(hex_values) >= num_traces * block_size:
                    plaintexts = np.zeros((num_traces, block_size), dtype=np.uint8)
                    for i in range(num_traces):
                        for j in range(block_size):
                            idx = i * block_size + j
                            if idx < len(hex_values):
                                plaintexts[i, j] = int(hex_values[idx], 16)
                    
                    print(f"  FromHeader fileLoadPlaintext: {num_traces} × {block_size} 字节")
                    return plaintexts
        except Exception as e:
            print(f"  Warning: LoadPlaintextFailed: {e}")
        
        return None
    
    def run(self, num_traces: int = 100, key: Optional[str] = None, 
            delay_us: int = 1000, output_file: Optional[str] = None,
            include_key: bool = True, include_plaintexts: bool = True):
        """RunComplete的采集流程（Use固定操作）"""
        self.traces = []
        
        if output_file is None:
            output_file = f"{self.algorithm}_traces.npz"
        
        # 1. GenerateTestData
        if not self.generate_test_data(num_traces, key, delay_us):
            return False
        
        # 2. 编译固件
        if not self._compile_firmware():
            return False
        
        # 3. 采集Trace
        if not self.run_qemu_and_collect():
            return False
        
        # 4. SaveData
        if not self.save_traces(output_file, include_key, include_plaintexts):
            return False
        
        print(f"\n{'='*70}")
        print(f"  {self.config['description']} Trace采集Completed！")
        print(f"  操作Class型: {self.config['operation']}")
        print(f"  Output file: {output_file}")
        print(f"{'='*70}")
        
        return True

def main():
    parser = argparse.ArgumentParser(
        description='STM32 Firmware Side-Channel AttackTrace采集器（简化版 - 每种Algorithm固定一种操作）',
        epilog='''
每种Algorithm的固定操作:
  AES:        encrypt  (S-boxOutput泄露)
  SM4:        encrypt  (S-boxOutput泄露)
  SM2:        sign     (Scalar乘泄露)
  RSA:        decrypt  (模幂运算泄露)
  ECC:        sign     (Scalar乘泄露)
  Dilithium:  sign     (NTT变换泄露)
        ''',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument('--algorithm', '-a', type=str, required=True,
                       choices=['aes', 'sm4', 'sm2', 'rsa', 'ecc', 'dilithium'],
                       help='AlgorithmClass型')
    
    parser.add_argument('--traces', '-n', type=int, default=100,
                       help='采集Trace数量 (Default: 100)')
    
    parser.add_argument('--key', type=str, default=None,
                       help='Key(hex)，Not指定则随机Generate')
    
    parser.add_argument('--delay', type=int, default=100000,
                       help='TRACE_START前的延时(微秒) (Default: 100000, i.e.100ms)')
    
    parser.add_argument('--output', '-o', type=str, default=None,
                       help='Output file名 (Default: {algorithm}_traces.npz)')
    
    parser.add_argument('--work-dir', type=str, default=None,
                       help='Working directory (Default: stm32_{algorithm}/)')
    
    parser.add_argument('--no-key', action='store_true',
                       help='NotInNPZ中SaveKey')
    
    parser.add_argument('--no-plaintexts', action='store_true',
                       help='NotInNPZ中SavePlaintext')
    
    args = parser.parse_args()
    
    try:
        collector = FirmwareTraceCollector(
            algorithm=args.algorithm,
            work_dir=args.work_dir
        )
        
        success = collector.run(
            num_traces=args.traces,
            key=args.key,
            delay_us=args.delay,
            output_file=args.output,
            include_key=not args.no_key,
            include_plaintexts=not args.no_plaintexts
        )
        
        return 0 if success else 1
        
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == '__main__':
    exit(main())
